package cn.codelead.masterdb.controller;

import cn.codelead.masterdb.config.LoggingAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@RestController
@CrossOrigin
public class OpenAiController {

    private final ChatClient chatClient;

    public OpenAiController(ChatClient.Builder chatClientBuilder,
                            VectorStore vectorStore,
                            ChatMemory chatMemory,
                            // mcp tools
                            ToolCallbackProvider mcpTools) {
        this.chatClient = chatClientBuilder
                .defaultSystem("""
                                       You are **DatabaseMaster**, an AI specialized in MySQL database analysis and querying. \s
                                                   Your responsibilities:
                                                   1. **Intent Understanding** \s
                                                      - 从用户对话中准确识别意图（查询、统计、结构分析等）。 \s
                                                      - 如有歧义，主动提问澄清。
                                                      - 注意，在执行sql之前你需要明确了解相应的表结构，避免问答中出现假设字段等等。如遇到不存在对应的业务、表等，也需明确向客户指出
                                                   
                                                   2. **Schema-Aware SQL Generation** \s
                                                      - 根据数据库模式（表名、字段、索引等）和用户意图，生成高效、可执行的 SQL 语句。 \s
                                                      - 尽量避免全表扫描，优先使用索引字段。
                                                   
                                                   3. **MCP 调用与结果处理** \s
                                                      - 通过 MCP 工具执行 SQL。 \s
                                                      - 捕获异常并返回友好错误提示。
                                                   
                                                   4. **Markdown 格式化输出** \s
                                                      - 始终以 Markdown 响应： \s
                                                        - **一级/二级/三级标题**（`#`,`##`,`###`） \s
                                                        - **表格**（`| 列1 | 列2 |`）展示结果集 \s
                                                        - **代码块**（````sql ... ````）展示 SQL \s
                                                        - **图片**（`![说明](url)`）或流程图（Mermaid）解释复杂逻辑 \s
                                                      - 在表格下方提供简要解读。
                                                   
                                                   5. **示例交互** \s
                                                      - 用户：查询过去 7 天的订单总额。 \s
                                                      - 你： \s
                                                        ```sql
                                                        SELECT DATE(order_date) AS date,
                                                               SUM(amount) AS total_amount
                                                        FROM orders
                                                        WHERE order_date >= CURDATE() - INTERVAL 7 DAY
                                                        GROUP BY DATE(order_date)
                                                        ```
                                                        | date       | total_amount |
                                                        |------------|--------------|
                                                        | 2025-04-04 | 12345.67     |
                                                        | …          | …            |
                                                        **解读**：过去 7 天每日订单总额如上。
                                                   
                                                   如果用户需要更多信息（如按渠道、按地区分组），请给出可选 SQL 并说明如何修改。 \s
                                                   
                        """) // 系统提示词
                .defaultAdvisors(
                        new PromptChatMemoryAdvisor(chatMemory),
                        new LoggingAdvisor())
//                .defaultTools(anytools)
                .defaultTools(mcpTools)
                //.defaultFunctions("func1", "func2", "func3") // FUNCTION CALLING
                .build();
    }

    @Autowired
    private VectorStore vectorStore;

    @CrossOrigin
    @GetMapping(value = "/ai/generateStreamAsString", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStreamAsString(@RequestParam(value = "message", defaultValue = "") String message) {

        Flux<String> content = chatClient.prompt()
                .system(s -> s.param("current_date", LocalDate.now().toString()))
                .advisors(a -> a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .user(message)
                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().query(message).build()))
                .stream()
                .content();

        return content
                .concatWith(Flux.just("[complete]"));

    }


}
